#include <cmath>
#include <fstream>
#include <iostream>

#include "gtsam/base/Lie.h"
#include "gtsam/base/Matrix.h"
#include "gtsam/base/Vector.h"

// ref: https://gtsam.org/notes/GTSAM-Concepts.html
// ref: https://dongjing3309.github.io/files/gtsam-tutorial.pdf
struct Point2c {
  double x;
  double y;
  Point2c(double xx, double yy) : x(xx), y(yy) {}
};

/**
 * any type compatible with GTSAM must have (in its traits):
 * - print function with an optional begining string, in format:
 *   static void Print(const T& m1, const std::string& str = "");
 * - Equal function with optional toterance, in format:
 *   static bool Equals(const T& m1, const T& m2, double tol = 1e-8);
 *
 * A mainfold type must have:
 * - int dimension (not fully support dynamic dimensionality yet)
 * - Typedefs TangentVector, where TangentVector = Eigen::Matrix<double,
 * dimension, 1>
 * - Local coordiante function, in format:
 *   static TangentVector Local(const Class& origin, const Class& other);
 * - Retraction back to manifold, in format:
 *   static Class Retract(const Class& origin, const TangentVector& v)
 *
 * A lie group types must have:
 * - Identity function
 * - Logmap function, with optional jacobians
 * - Expmap function, with optional jacobians
 * - Compose function, with optional jacobians
 * - Between function, with optional jacobians
 * - Inverse function, with optional jacobians
 *
 * For lie group types, other than traits, operator * or (+ and -) should be
 * defined for compose / between operation. can be defined inside or outside
 * class.
 */

// traits must in namespace gtsam
namespace gtsam {

template <>
struct traits<Point2c> {
  // structural catagory: this is a lie group
  // availble options: manifold_tag, group_tag, lie_group_tag
  typedef lie_group_tag structure_category;

  /**
   * Basic (Testable)
   */
  // print
  static void Print(const Point2c& m, const std::string str = "") {
    std::cout << str << "(" << m.x << ", " << m.y << ")" << std::endl;
  }

  // equality with optional tol
  static bool Equals(const Point2c& m1,
                     const Point2c& m2,
                     const double tol = 1e-8) {
    if (std::fabs(m1.x - m2.x) < tol && std::fabs(m1.y - m2.y) < tol) {
      return true;
    } else {
      return false;
    }
  }

  /**
   * Manifold
   */
  // use enum dimension
  enum { dimension = 2 };
  static int GetDimension(const Point2c&) {
    return dimension;
  }

  // Typedefs needed
  typedef Point2c ManifoldType;
  typedef Eigen::Matrix<double, dimension, 1> TangentVector;

  // Local coordinate of Point2c is native (vectorspace)
  static TangentVector Local(const Point2c& origin, const Point2c& other) {
    return TangentVector(other.x - origin.x, other.y - origin.x);
  }

  // Retraction back to mainfold of Point2c is native (vectorspace)
  static Point2c Retract(const Point2c& origin, const TangentVector& v) {
    return Point2c(origin.x + v[0], origin.y + v[1]);
  }

  /**
   * Lie Group
   */
  // indicate this group using operation *,
  // if use +/- then use option additive_group_tag
  typedef multiplicative_group_tag group_flavor;

  // typdefs
  typedef OptionalJacobian<dimension, dimension> ChartJacobian;

  static Point2c Identity() {
    return Point2c(0, 0);
  }

  static TangentVector Logmap(const Point2c& m,
                              ChartJacobian Hm = boost::none) {
    if (Hm) {
      (*Hm) = Matrix2::Identity();
    }
    return TangentVector(m.x, m.y);
  }

  static Point2c Expmap(const TangentVector& v,
                        ChartJacobian Hv = boost::none) {
    if (Hv) {
      (*Hv) = Matrix2::Identity();
    }
    return Point2c(v(0), v(1));
  }

  static Point2c Compose(const Point2c& m1,
                         const Point2c& m2,
                         ChartJacobian H1 = boost::none,
                         ChartJacobian H2 = boost::none) {
    if (H1) {
      (*H1) = Matrix2::Identity();
    }
    if (H2) {
      (*H2) = Matrix2::Identity();
    }
    return Point2c(m1.x + m2.x, m1.y + m2.y);
  }

  static Point2c Between(const Point2c& m1,
                         const Point2c& m2,
                         ChartJacobian H1 = boost::none,
                         ChartJacobian H2 = boost::none) {
    if (H1) {
      (*H1) = -Matrix2::Identity();
    }
    if (H2) {
      (*H2) = Matrix2::Identity();
    }
    return Point2c(m2.x - m1.x, m2.y - m1.y);
  }

  static Point2c Inverse(const Point2c& m, ChartJacobian H = boost::none) {
    if (H) {
      *H = -Matrix2::Identity();
    }
    return Point2c(-m.x, -m.y);
  }
};

}  // namespace gtsam

// operator *
Point2c operator*(const Point2c& m1, const Point2c& m2) {
  return Point2c(m1.x + m2.x, m1.y + m2.y);
}

#include <gtsam/inference/Symbol.h>
#include <gtsam/nonlinear/GaussNewtonOptimizer.h>
#include <gtsam/nonlinear/Marginals.h>
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam/nonlinear/Values.h>
#include <gtsam/slam/BetweenFactor.h>
#include <gtsam/slam/PriorFactor.h>

using namespace std;
using namespace gtsam;

int main(int argc, char** argv) {
  gtsam::NonlinearFactorGraph graph;
  noiseModel::Diagonal::shared_ptr priorModel =
      noiseModel::Diagonal::Sigmas(Vector2(0.2, 0.2));
  graph.add(PriorFactor<Point2c>(Symbol('x', 1), Point2c(0, 0), priorModel));
  noiseModel::Diagonal::shared_ptr odomModel =
      noiseModel::Diagonal::Sigmas(Vector2(0.5, 0.5));

  graph.add(BetweenFactor<Point2c>(
      Symbol('x', 1), Symbol('x', 2), Point2c(2, 0), odomModel));
  graph.add(BetweenFactor<Point2c>(
      Symbol('x', 2), Symbol('x', 3), Point2c(2, 0), odomModel));
  graph.add(BetweenFactor<Point2c>(
      Symbol('x', 3), Symbol('x', 4), Point2c(2, 0), odomModel));
  graph.add(BetweenFactor<Point2c>(
      Symbol('x', 4), Symbol('x', 5), Point2c(2, 0), odomModel));

  graph.print("\nFactor Graph:\n");

  Values initials;
  initials.insert(Symbol('x', 1), Point2c(0.2, -0.3));
  initials.insert(Symbol('x', 2), Point2c(2.1, 0.3));
  initials.insert(Symbol('x', 3), Point2c(3.9, -0.1));
  initials.insert(Symbol('x', 4), Point2c(5.9, -0.3));
  initials.insert(Symbol('x', 5), Point2c(8.2, 0.1));

  initials.print("\nInitial Values:\n");

  GaussNewtonParams parameters;
  parameters.setVerbosity("ERROR");

  GaussNewtonOptimizer optimizer(graph, initials, parameters);
  Values results = optimizer.optimize();

  results.print("Final Result:\n");

  Marginals marginals(graph, results);

  cout << "x1 covariance:\n"
       << marginals.marginalCovariance(Symbol('x', 1)) << endl;
  cout << "x2 covariance:\n"
       << marginals.marginalCovariance(Symbol('x', 2)) << endl;
  cout << "x3 covariance:\n"
       << marginals.marginalCovariance(Symbol('x', 3)) << endl;
  cout << "x4 covariance:\n"
       << marginals.marginalCovariance(Symbol('x', 4)) << endl;
  cout << "x5 covariance:\n"
       << marginals.marginalCovariance(Symbol('x', 5)) << endl;
}